// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "fastflow/framework/execution_provider.h"

#include "fastflow/graph/graph.h"
#include "fastflow/framework/computation_capacity.h"
#include "fastflow/framework/kernel_registry_manager.h"
#include "fastflow/framework/op_kernel.h"
#include "fastflow/framework/kernel_registry.h"
#include "fastflow/graph/graph_viewer.h"

namespace fastflow {

namespace {
inline int MakeKey(int id, FASTFLOWMemType mem_type) {
  return id << 2 | mem_type;
}
}  // namespace

AllocatorPtr IExecutionProvider::GetAllocator(int id, FASTFLOWMemType mem_type) const {
  auto iter = allocators_.find(MakeKey(id, mem_type));
  if (iter != allocators_.end()) {
    return iter->second;
  }
  return nullptr;
}

std::vector<std::unique_ptr<ComputationCapacity>>
IExecutionProvider::GetCapability(const fastflow::GraphViewer& graph,
                                 const std::vector<const KernelRegistry*>& kernel_registries) const {
 std::vector<std::unique_ptr<ComputationCapacity>> result;
 for (auto& node : graph.Nodes()) {
   for (auto registry : kernel_registries) {
     if (registry->TryFindKernel(node, Type()) != nullptr) {
       std::unique_ptr<IndexedSubGraph> sub_graph = std::make_unique<IndexedSubGraph>();
       sub_graph->nodes.push_back(node.Index());
       result.push_back(std::make_unique<ComputationCapacity>(std::move(sub_graph), nullptr));
     }
   }
 }

 return result;
}

Status IExecutionProvider::CopyTensor(const Tensor& src,
                                              Tensor& dst,
                                              int exec_queue_id) const {
  // execution provider may override this to support different exec queues
  FASTFLOW_ENFORCE(exec_queue_id == 0);
  return CopyTensor(src, dst);
}

Status IExecutionProvider::Sync() const { return Status::OK(); };

Status IExecutionProvider::OnRunStart() { return Status::OK(); }

Status IExecutionProvider::OnRunEnd() { return Status::OK(); }

void IExecutionProvider::InsertAllocator(AllocatorPtr allocator) {
  const FASTFLOWAllocatorInfo& info = allocator->Info();
  const int key = MakeKey(info.id, info.mem_type);
  auto iter = allocators_.find(key);
  if (iter != allocators_.end()) {
    FASTFLOW_THROW("duplicated allocator");
  }
  allocators_.insert(iter, {key, allocator});
}
}  // namespace fastflow
